package com.shujia.mllib

import org.apache.spark.ml.classification.{NaiveBayes, NaiveBayesModel}
import org.apache.spark.ml.feature.{HashingTF, IDF, IDFModel, Tokenizer}
import org.apache.spark.sql.expressions.UserDefinedFunction
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

object Demo6TestClass {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession
      .builder()
      .master("local[8]")
      .appName("image")
      .getOrCreate()
    import spark.implicits._
    import org.apache.spark.sql.functions._

    //1、读取原始数据
    val linesDF: DataFrame = spark.read
      .format("csv")
      .option("sep", "\t")
      .schema("label DOUBLE,text STRING")
      .load("spark/data/text.txt")
      .repartition(8)

    //分词自定义函数
    val ikUDF: UserDefinedFunction = udf((text: String) => {
      val words: List[String] = IKUtil.fit(text)
      //将词语使用空格拼接返回
      words.mkString(" ")
    })

    val wordsDF: Dataset[Row] = linesDF
      //使用ik分词器分词
      .select($"label", ikUDF($"text") as "sentence")
      //过滤空数据
      .filter($"sentence" =!= "")


    //英文分词器
    val tokenizer: Tokenizer = new Tokenizer()
      .setInputCol("sentence") //输入列名
      .setOutputCol("words") //输出列名

    val wordsData: DataFrame = tokenizer.transform(wordsDF)


    //增加TF(构建词典向量)
    val hashingTF: HashingTF = new HashingTF()
      .setInputCol("words") //输入列名
      .setOutputCol("rawFeatures") //输出列名
    //.setNumFeatures(20) //总的单词的数量

    val featurizedData: DataFrame = hashingTF.transform(wordsData)

    //逆文本频率
    val idf: IDF = new IDF()
      .setInputCol("rawFeatures")
      .setOutputCol("features")

    //训练逆文本频率的模型
    val idfModel: IDFModel = idf.fit(featurizedData)

    val rescaledData: DataFrame = idfModel
      .transform(featurizedData)


    /**
     * 构建算法，训练模型，测试模型的准确率
     *
     */

    //将数据拆分成训练集和测试集‘
    val Array(train: DataFrame, test: DataFrame) = rescaledData.randomSplit(Array(0.7, 0.3))

    //文本分类可以使用贝叶斯分类
    val naiveBayes = new NaiveBayes()

    //将训练集带入算法训练模型
    val model: NaiveBayesModel = naiveBayes.fit(train)

    //将测试集带入模型测试模型的准确率
    val frame: DataFrame = model.transform(test)

    frame.show(false)
    //计算准确率
    //6、计算准确率
    val p: Double = frame.where($"label" === $"prediction").count() / frame.count().toDouble
    println(s"模型的准确率：${p}")
  }

}
